{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KesQRvxkm1PH"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "PROJECT = \"PROJECT\" # REPLACE WITH YOUR PROJECT ID\n",
    "BUCKET = \"BUCKET\" # REPLACE WITH A BUCKET NAME (PUT YOUR PROJECT ID AND WE CREATE THE BUCKET ITSELF NEXT)\n",
    "REGION = \"us-central1\" # REPLACE WITH YOUR REGION e.g. us-central1\n",
    "\n",
    "# do not change these\n",
    "os.environ[\"PROJECT\"] = PROJECT\n",
    "os.environ[\"BUCKET\"] =  BUCKET\n",
    "os.environ[\"REGION\"] = REGION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TLTq9A_dm1PL",
    "outputId": "5619ad91-841f-41ee-dc52-fd248fa2ad1f"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "gcloud config set project $PROJECT\n",
    "gcloud config set compute/region $REGION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "L4MZZN-Om1PQ"
   },
   "outputs": [],
   "source": [
    "# %%bash\n",
    "# OUTDIR=gs://${BUCKET}/taxifare/trained_model/yellow_trips\n",
    "# JOBNAME=taxifare_yellow_trips_$(date -u +%y%m%d_%H%M%S)\n",
    "# echo $OUTDIR $REGION $JOBNAME\n",
    "# # Clear the Cloud Storage Bucket used for the training job\n",
    "# gsutil -m rm -rf $OUTDIR\n",
    "# gcloud ml-engine jobs submit training $JOBNAME \\\n",
    "#   --region=$REGION \\\n",
    "#   --module-name=trainer.task \\\n",
    "#   --package-path=${PWD}/cloud_composer_automated_ml_pipeline_taxifare_module/trainer \\\n",
    "#   --job-dir=$OUTDIR \\\n",
    "#   --staging-bucket=gs://$BUCKET \\\n",
    "#   --scale-tier=BASIC \\\n",
    "#   --runtime-version=\"1.13\" \\\n",
    "#   -- \\\n",
    "#   --train_data_paths=gs://$BUCKET/taxifare/data/yellow_trips/train-* \\\n",
    "#   --eval_data_paths=gs://$BUCKET/taxifare/data/yellow_trips/valid-* \\\n",
    "#   --output_dir=$OUTDIR \\\n",
    "#   --train_steps=500"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HxR6rmGBm1PT"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "cd cloud_composer_automated_ml_pipeline_taxifare_module\n",
    "touch README.md\n",
    "python setup.py sdist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fjzhe-U2m1PV"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "gsutil cp cloud_composer_automated_ml_pipeline_taxifare_module/dist/taxifare-0.1.tar.gz gs://$BUCKET/taxifare/code/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "O241YkN7m1PX"
   },
   "source": [
    "***\n",
    "# Part Two: Setup a scheduled workflow with Cloud Composer\n",
    "In this section you will complete a partially written training.py DAG file and copy it to the DAGS folder in your Composer instance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BczwG3zsm1PY"
   },
   "source": [
    "## Copy your Airflow bucket name\n",
    "1. Navigate to your Cloud Composer [instance](https://console.cloud.google.com/composer/environments?project=)<br/><br/>\n",
    "2. Select __DAGs Folder__<br/><br/>\n",
    "3. You will be taken to the Google Cloud Storage bucket that Cloud Composer has created automatically for your Airflow instance<br/><br/>\n",
    "4. __Copy the bucket name__ into the variable below (example: us-central1-composer-08f6edeb-bucket)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1r6tYEaDm1PZ"
   },
   "outputs": [],
   "source": [
    "AIRFLOW_BUCKET = \"AIRFLOW_BUCKET\" # REPLACE WITH AIRFLOW BUCKET NAME\n",
    "os.environ[\"AIRFLOW_BUCKET\"] = AIRFLOW_BUCKET"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7viNmxubm1Pb"
   },
   "source": [
    "## Complete the training.py DAG file\n",
    "Apache Airflow orchestrates tasks out to other services through a [DAG (Directed Acyclic Graph)](https://airflow.apache.org/concepts.html) file which specifies what services to call, what to do, and when to run these tasks. DAG files are written in python and are loaded automatically into Airflow once present in the Airflow/dags/ folder in your Cloud Composer bucket. \n",
    "\n",
    "Execute the code cells to create the files."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile airflow/dags/taxifare_multi.py\n",
    "# Copyright 2018 Google Inc. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "\n",
    "\"\"\"DAG definition for taxifare automated pipeline.\"\"\"\n",
    "\n",
    "import airflow\n",
    "from airflow import DAG\n",
    "\n",
    "# Reference for all available airflow operators: \n",
    "# https://github.com/apache/incubator-airflow/tree/master/airflow/contrib/operators\n",
    "from airflow.contrib.operators.bigquery_check_operator import BigQueryCheckOperator\n",
    "from airflow.contrib.operators.bigquery_operator import BigQueryOperator\n",
    "from airflow.contrib.operators.bigquery_to_gcs import BigQueryToCloudStorageOperator\n",
    "from airflow.operators.bash_operator import BashOperator\n",
    "from airflow.operators.python_operator import BranchPythonOperator\n",
    "from airflow.operators.dummy_operator import DummyOperator\n",
    "from airflow.hooks.base_hook import BaseHook\n",
    "\n",
    "from airflow.contrib.operators.mlengine_operator import MLEngineTrainingOperator, MLEngineModelOperator, MLEngineVersionOperator\n",
    "from airflow.models import TaskInstance\n",
    "\n",
    "import datetime\n",
    "import logging\n",
    "\n",
    "def _get_project_id():\n",
    "  \"\"\"Get project ID from default GCP connection.\"\"\"\n",
    "\n",
    "  extras = BaseHook.get_connection(\"google_cloud_default\").extra_dejson\n",
    "  key = \"extra__google_cloud_platform__project\"\n",
    "  if key in extras:\n",
    "    project_id = extras[key]\n",
    "  else:\n",
    "    raise (\"Must configure project_id in google_cloud_default \"\n",
    "           \"connection from Airflow Console\")\n",
    "  return project_id\n",
    "\n",
    "PROJECT_ID = _get_project_id()\n",
    "\n",
    "# Data set constants, used in BigQuery tasks.    You can change these\n",
    "# to conform to your data.\n",
    "\n",
    "# Specify your source BigQuery project, dataset, and table names\n",
    "SOURCE_BQ_PROJECT = \"nyc-tlc\"\n",
    "SOURCE_DATASET_TABLE_NAMES = \"yellow.trips,green.trips_2014,green.trips_2015\".split(\",\")\n",
    "\n",
    "# Specify your destination BigQuery dataset\n",
    "DESTINATION_DATASET = \"taxifare\"\n",
    "\n",
    "# GCS bucket names and region, can also be changed.\n",
    "BUCKET = \"gs://\" + PROJECT_ID + \"-bucket\"\n",
    "REGION = \"us-east1\"\n",
    "\n",
    "# # The code package name comes from the model code in the wals_ml_engine\n",
    "# # directory of the solution code base.\n",
    "PACKAGE_URI = BUCKET + \"/taxifare/code/taxifare-0.1.tar.gz\"\n",
    "JOB_DIR = BUCKET + \"/jobs\"\n",
    "\n",
    "default_args = {\n",
    "  \"owner\": \"airflow\",\n",
    "  \"depends_on_past\": False,\n",
    "  \"start_date\": airflow.utils.dates.days_ago(2),\n",
    "  \"email\": [\"airflow@example.com\"],\n",
    "  \"email_on_failure\": True,\n",
    "  \"email_on_retry\": False,\n",
    "  \"retries\": 5,\n",
    "  \"retry_delay\": datetime.timedelta(minutes=5)\n",
    "}\n",
    "\n",
    "# Default schedule interval using cronjob syntax - can be customized here\n",
    "# or in the Airflow console.\n",
    "\n",
    "# Specify a schedule interval in CRON syntax to run once a day at 2100 hours (9pm)\n",
    "# Reference: https://airflow.apache.org/scheduler.html\n",
    "schedule_interval = \"00 21 * * *\"\n",
    "\n",
    "# Title your DAG\n",
    "dag = DAG(\n",
    "  \"taxifare_multi\", \n",
    "  default_args=default_args,\n",
    "  schedule_interval=None\n",
    ")\n",
    "\n",
    "dag.doc_md = __doc__\n",
    "\n",
    "\n",
    "#\n",
    "#\n",
    "# Task Definition\n",
    "#\n",
    "#\n",
    "\n",
    "for model in SOURCE_DATASET_TABLE_NAMES:\n",
    "  # BigQuery data query\n",
    "  bql=\"\"\"\n",
    "  SELECT\n",
    "    (tolls_amount + fare_amount) AS fare_amount,\n",
    "    EXTRACT(DAYOFWEEK FROM pickup_datetime) * 1.0 AS dayofweek,\n",
    "    EXTRACT(HOUR FROM pickup_datetime) * 1.0 AS hourofday,\n",
    "    pickup_longitude AS pickuplon,\n",
    "    pickup_latitude AS pickuplat,\n",
    "    dropoff_longitude AS dropofflon,\n",
    "    dropoff_latitude AS dropofflat,\n",
    "    passenger_count*1.0 AS passengers,\n",
    "    CONCAT(CAST(pickup_datetime AS STRING), CAST(pickup_longitude AS STRING), CAST(pickup_latitude AS STRING), CAST(dropoff_latitude AS STRING), CAST(dropoff_longitude AS STRING)) AS key\n",
    "  FROM\n",
    "    `{0}.{1}`\n",
    "  WHERE\n",
    "    trip_distance > 0\n",
    "    AND fare_amount >= 2.5\n",
    "    AND pickup_longitude > -78\n",
    "    AND pickup_longitude < -70\n",
    "    AND dropoff_longitude > -78\n",
    "    AND dropoff_longitude < -70\n",
    "    AND pickup_latitude > 37\n",
    "    AND pickup_latitude < 45\n",
    "    AND dropoff_latitude > 37\n",
    "    AND dropoff_latitude < 45\n",
    "    AND passenger_count > 0\n",
    "    AND rand() < 0.00001\n",
    "  \"\"\"\n",
    "\n",
    "  bql = bql.format(SOURCE_BQ_PROJECT, model)\n",
    "\n",
    "  bql_train = \"SELECT * EXCEPT (key) FROM({0}) WHERE ABS(MOD(FARM_FINGERPRINT(key), 5)) < 4\".format(bql)\n",
    "  bql_eval = \"SELECT * EXCEPT (key) FROM({0}) WHERE ABS(MOD(FARM_FINGERPRINT(key), 5)) = 4\".format(bql)\n",
    "\n",
    "  # Complete the BigQueryOperator task to truncate the table if it already exists before writing\n",
    "  # Reference: https://airflow.apache.org/integration.html#bigqueryoperator\n",
    "  bq_train_data_op = BigQueryOperator(\n",
    "    task_id=\"bq_train_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bql=bql_train,\n",
    "    destination_dataset_table=\"{}.{}_train_data\".format(DESTINATION_DATASET, model.replace(\".\",\"_\")),\n",
    "    write_disposition=\"WRITE_TRUNCATE\", # specify to truncate on writes\n",
    "    use_legacy_sql=False,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bq_eval_data_op = BigQueryOperator(\n",
    "    task_id=\"bq_eval_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bql=bql_eval,\n",
    "    destination_dataset_table=\"{}.{}_eval_data\".format(DESTINATION_DATASET, model.replace(\".\",\"_\")),\n",
    "    write_disposition=\"WRITE_TRUNCATE\", # specify to truncate on writes\n",
    "    use_legacy_sql=False,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  sql = \"\"\"\n",
    "  SELECT\n",
    "    COUNT(*)\n",
    "  FROM\n",
    "    [{0}:{1}.{2}]\n",
    "  \"\"\"\n",
    "\n",
    "  # Check to make sure that the data tables won\"t be empty\n",
    "  bq_check_train_data_op = BigQueryCheckOperator(\n",
    "    task_id=\"bq_check_train_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    sql=sql.format(PROJECT_ID, DESTINATION_DATASET, model.replace(\".\",\"_\") + \"_train_data\"),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bq_check_eval_data_op = BigQueryCheckOperator(\n",
    "    task_id=\"bq_check_eval_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    sql=sql.format(PROJECT_ID, DESTINATION_DATASET, model.replace(\".\",\"_\") + \"_eval_data\"),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # BigQuery training data export to GCS\n",
    "  bash_remove_old_data_op = BashOperator(\n",
    "    task_id=\"bash_remove_old_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bash_command=\"if gsutil ls {0}/taxifare/data/{1} 2> /dev/null; then gsutil -m rm -rf {0}/taxifare/data/{1}/*; else true; fi\".format(BUCKET, model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # Takes a BigQuery dataset and table as input and exports it to GCS as a CSV\n",
    "  train_files = BUCKET + \"/taxifare/data/\"\n",
    "\n",
    "  bq_export_gcs_train_csv_op = BigQueryToCloudStorageOperator(\n",
    "    task_id=\"bq_export_gcs_train_csv_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    source_project_dataset_table=\"{}.{}_train_data\".format(DESTINATION_DATASET, model.replace(\".\",\"_\")),\n",
    "    destination_cloud_storage_uris=[train_files + \"{}/train-*.csv\".format(model.replace(\".\",\"_\"))],\n",
    "    export_format=\"CSV\",\n",
    "    print_header=False,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  eval_files = BUCKET + \"/taxifare/data/\"\n",
    "\n",
    "  bq_export_gcs_eval_csv_op = BigQueryToCloudStorageOperator(\n",
    "    task_id=\"bq_export_gcs_eval_csv_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    source_project_dataset_table=\"{}.{}_eval_data\".format(DESTINATION_DATASET, model.replace(\".\",\"_\")),\n",
    "    destination_cloud_storage_uris=[eval_files + \"{}/eval-*.csv\".format(model.replace(\".\",\"_\"))],\n",
    "    export_format=\"CSV\",\n",
    "    print_header=False,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "\n",
    "  # ML Engine training job\n",
    "  job_id = \"taxifare_{}_{}\".format(model.replace(\".\",\"_\"), datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\"))\n",
    "  output_dir = BUCKET + \"/taxifare/trained_model/{}\".format(model.replace(\".\",\"_\"))\n",
    "  job_dir = JOB_DIR + \"/\" + job_id\n",
    "  training_args = [\n",
    "    \"--job-dir\", job_dir,\n",
    "    \"--train_data_paths\", train_files,\n",
    "    \"--eval_data_paths\", eval_files,\n",
    "    \"--output_dir\", output_dir,\n",
    "    \"--train_steps\", str(500),\n",
    "    \"--train_batch_size\", str(32),\n",
    "    \"--eval_steps\", str(500),\n",
    "    \"--eval_batch_size\", str(32),\n",
    "    \"--nbuckets\", str(8),\n",
    "    \"--hidden_units\", \"128,32,4\"\n",
    "  ]\n",
    "\n",
    "  # Reference: https://airflow.apache.org/integration.html#cloud-ml-engine\n",
    "  ml_engine_training_op = MLEngineTrainingOperator(\n",
    "    task_id=\"ml_engine_training_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID,\n",
    "    job_id=job_id,\n",
    "    package_uris=[PACKAGE_URI],\n",
    "    training_python_module=\"trainer.task\",\n",
    "    training_args=training_args,\n",
    "    region=REGION,\n",
    "    scale_tier=\"BASIC\",\n",
    "    runtime_version=\"1.13\", \n",
    "    python_version=\"3.5\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  MODEL_NAME = \"taxifare_\"\n",
    "  MODEL_VERSION = \"v1\"\n",
    "  MODEL_LOCATION = BUCKET + \"/taxifare/saved_model/\"\n",
    "\n",
    "  bash_remove_old_saved_model_op = BashOperator(\n",
    "    task_id=\"bash_remove_old_saved_model_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bash_command=\"if gsutil ls {0} 2> /dev/null; then gsutil -m rm -rf {0}/*; else true; fi\".format(MODEL_LOCATION + model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bash_copy_new_saved_model_op = BashOperator(\n",
    "    task_id=\"bash_copy_new_saved_model_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bash_command=\"gsutil -m rsync -d -r `gsutil ls {0}/export/exporter/ | tail -1` {1}\".format(output_dir, MODEL_LOCATION + model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # Create model on ML-Engine\n",
    "  bash_ml_engine_models_list_op = BashOperator(\n",
    "    task_id=\"bash_ml_engine_models_list_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    xcom_push=True,\n",
    "    bash_command=\"gcloud ml-engine models list --filter='name:{0}'\".format(MODEL_NAME + model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  def check_if_model_already_exists(templates_dict, **kwargs):\n",
    "    cur_model = templates_dict[\"model\"].replace(\".\",\"_\")\n",
    "    ml_engine_models_list = kwargs[\"ti\"].xcom_pull(task_ids=\"bash_ml_engine_models_list_{}_task\".format(cur_model))\n",
    "    logging.info(\"check_if_model_already_exists: {}: ml_engine_models_list = \\n{}\".format(cur_model, ml_engine_models_list))\n",
    "    create_model_task = \"ml_engine_create_model_{}_task\".format(cur_model)\n",
    "    dont_create_model_task = \"dont_create_model_dummy_branch_{}_task\".format(cur_model)\n",
    "    if len(ml_engine_models_list) == 0 or ml_engine_models_list == \"Listed 0 items.\":\n",
    "      return create_model_task\n",
    "    return dont_create_model_task\n",
    "\n",
    "  check_if_model_already_exists_op = BranchPythonOperator(\n",
    "    task_id=\"check_if_model_already_exists_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    templates_dict={\"model\": model.replace(\".\",\"_\")},\n",
    "    python_callable=check_if_model_already_exists,\n",
    "    provide_context=True,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_create_model_op = MLEngineModelOperator(\n",
    "    task_id=\"ml_engine_create_model_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model={\"name\": MODEL_NAME + model.replace(\".\",\"_\")}, \n",
    "    operation=\"create\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  create_model_dummy_op = DummyOperator(\n",
    "    task_id=\"create_model_dummy_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    trigger_rule=\"all_done\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  dont_create_model_dummy_branch_op = DummyOperator(\n",
    "    task_id=\"dont_create_model_dummy_branch_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  dont_create_model_dummy_op = DummyOperator(\n",
    "    task_id=\"dont_create_model_dummy_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    trigger_rule=\"all_done\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # Create version of model on ML-Engine\n",
    "  bash_ml_engine_versions_list_op = BashOperator(\n",
    "    task_id=\"bash_ml_engine_versions_list_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    xcom_push=True,\n",
    "    bash_command=\"gcloud ml-engine versions list --model {0} --filter='name:{1}'\".format(MODEL_NAME + model.replace(\".\",\"_\"), MODEL_VERSION),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  def check_if_model_version_already_exists(templates_dict, **kwargs):\n",
    "    cur_model = templates_dict[\"model\"].replace(\".\",\"_\")\n",
    "    ml_engine_versions_list = kwargs[\"ti\"].xcom_pull(task_ids=\"bash_ml_engine_versions_list_{}_task\".format(cur_model))\n",
    "    logging.info(\"check_if_model_version_already_exists: {}: ml_engine_versions_list = \\n{}\".format(cur_model, ml_engine_versions_list))\n",
    "    create_version_task = \"ml_engine_create_version_{}_task\".format(cur_model)\n",
    "    create_other_version_task = \"ml_engine_create_other_version_{}_task\".format(cur_model)\n",
    "    if len(ml_engine_versions_list) == 0 or ml_engine_versions_list == \"Listed 0 items.\":\n",
    "      return create_version_task\n",
    "    return create_other_version_task\n",
    "\n",
    "  check_if_model_version_already_exists_op = BranchPythonOperator(\n",
    "    task_id=\"check_if_model_version_already_exists_{}_task\".format(model.replace(\".\",\"_\")), \n",
    "    templates_dict={\"model\": model.replace(\".\",\"_\")},\n",
    "    python_callable=check_if_model_version_already_exists,\n",
    "    provide_context=True,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  OTHER_VERSION_NAME = \"v_{0}\".format(datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")[0:12])\n",
    "\n",
    "  ml_engine_create_version_op = MLEngineVersionOperator(\n",
    "    task_id=\"ml_engine_create_version_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model_name=MODEL_NAME + model.replace(\".\",\"_\"), \n",
    "    version_name=MODEL_VERSION, \n",
    "    version={\n",
    "      \"name\": MODEL_VERSION,\n",
    "      \"deploymentUri\": MODEL_LOCATION + model.replace(\".\",\"_\"),\n",
    "      \"runtimeVersion\": \"1.13\",\n",
    "      \"framework\": \"TENSORFLOW\",\n",
    "      \"pythonVersion\": \"3.5\",\n",
    "    },\n",
    "    operation=\"create\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_create_other_version_op = MLEngineVersionOperator(\n",
    "    task_id=\"ml_engine_create_other_version_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model_name=MODEL_NAME + model.replace(\".\",\"_\"), \n",
    "    version_name=OTHER_VERSION_NAME, \n",
    "    version={\n",
    "      \"name\": OTHER_VERSION_NAME,\n",
    "      \"deploymentUri\": MODEL_LOCATION + model.replace(\".\",\"_\"),\n",
    "      \"runtimeVersion\": \"1.13\",\n",
    "      \"framework\": \"TENSORFLOW\",\n",
    "      \"pythonVersion\": \"3.5\",\n",
    "    },\n",
    "    operation=\"create\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_set_default_version_op = MLEngineVersionOperator(\n",
    "    task_id=\"ml_engine_set_default_version_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model_name=MODEL_NAME + model.replace(\".\",\"_\"), \n",
    "    version_name=MODEL_VERSION, \n",
    "    version={\"name\": MODEL_VERSION}, \n",
    "    operation=\"set_default\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_set_default_other_version_op = MLEngineVersionOperator(\n",
    "    task_id=\"ml_engine_set_default_other_version_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model_name=MODEL_NAME + model.replace(\".\",\"_\"), \n",
    "    version_name=OTHER_VERSION_NAME, \n",
    "    version={\"name\": OTHER_VERSION_NAME}, \n",
    "    operation=\"set_default\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # Build dependency graph, set_upstream dependencies for all tasks\n",
    "  bq_check_train_data_op.set_upstream(bq_train_data_op)\n",
    "  bq_check_eval_data_op.set_upstream(bq_eval_data_op)\n",
    "\n",
    "  bash_remove_old_data_op.set_upstream([bq_check_train_data_op, bq_check_eval_data_op])\n",
    "\n",
    "  bq_export_gcs_train_csv_op.set_upstream([bash_remove_old_data_op])\n",
    "  bq_export_gcs_eval_csv_op.set_upstream([bash_remove_old_data_op])\n",
    "\n",
    "  ml_engine_training_op.set_upstream([bq_export_gcs_train_csv_op, bq_export_gcs_eval_csv_op])\n",
    "\n",
    "  bash_remove_old_saved_model_op.set_upstream(ml_engine_training_op)\n",
    "  bash_copy_new_saved_model_op.set_upstream(bash_remove_old_saved_model_op)\n",
    "\n",
    "  bash_ml_engine_models_list_op.set_upstream(ml_engine_training_op)\n",
    "  check_if_model_already_exists_op.set_upstream(bash_ml_engine_models_list_op)\n",
    "\n",
    "  ml_engine_create_model_op.set_upstream(check_if_model_already_exists_op)\n",
    "  create_model_dummy_op.set_upstream(ml_engine_create_model_op)\n",
    "  dont_create_model_dummy_branch_op.set_upstream(check_if_model_already_exists_op)\n",
    "  dont_create_model_dummy_op.set_upstream(dont_create_model_dummy_branch_op)\n",
    "\n",
    "  bash_ml_engine_versions_list_op.set_upstream([dont_create_model_dummy_op, create_model_dummy_op])\n",
    "  check_if_model_version_already_exists_op.set_upstream(bash_ml_engine_versions_list_op)\n",
    "\n",
    "  ml_engine_create_version_op.set_upstream([bash_copy_new_saved_model_op, check_if_model_version_already_exists_op])\n",
    "  ml_engine_create_other_version_op.set_upstream([bash_copy_new_saved_model_op, check_if_model_version_already_exists_op])\n",
    "\n",
    "  ml_engine_set_default_version_op.set_upstream(ml_engine_create_version_op)\n",
    "  ml_engine_set_default_other_version_op.set_upstream(ml_engine_create_other_version_op)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile airflow/dags/module/taxifare_module.py\n",
    "# Copyright 2018 Google Inc. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "\n",
    "\"\"\"DAG definition for taxifare automated pipeline.\"\"\"\n",
    "\n",
    "import airflow\n",
    "from airflow import DAG\n",
    "\n",
    "# Reference for all available airflow operators: \n",
    "# https://github.com/apache/incubator-airflow/tree/master/airflow/contrib/operators\n",
    "from airflow.hooks.base_hook import BaseHook\n",
    "\n",
    "from airflow.models import TaskInstance\n",
    "\n",
    "import datetime\n",
    "\n",
    "from module import preprocess\n",
    "from module import training\n",
    "from module import deploy\n",
    "\n",
    "\n",
    "def _get_project_id():\n",
    "  \"\"\"Get project ID from default GCP connection.\"\"\"\n",
    "\n",
    "  extras = BaseHook.get_connection(\"google_cloud_default\").extra_dejson\n",
    "  key = \"extra__google_cloud_platform__project\"\n",
    "  if key in extras:\n",
    "    project_id = extras[key]\n",
    "  else:\n",
    "    raise (\"Must configure project_id in google_cloud_default \"\n",
    "           \"connection from Airflow Console\")\n",
    "  return project_id\n",
    "\n",
    "# Constants\n",
    "# Get project ID and GCS bucket\n",
    "PROJECT_ID = _get_project_id()\n",
    "BUCKET = \"gs://\" + PROJECT_ID + \"-bucket\"\n",
    "\n",
    "# Specify your source BigQuery dataset and table names\n",
    "SOURCE_DATASET_TABLE_NAMES = \"yellow.trips,green.trips_2014,green.trips_2015\".split(\",\")\n",
    "\n",
    "# Where to write out data in GCS\n",
    "DATA_DIR = BUCKET + \"/taxifare/data/\"\n",
    "\n",
    "# Base model parameters\n",
    "MODEL_NAME = \"taxifare_\"\n",
    "MODEL_VERSION = \"v1\"\n",
    "MODEL_LOCATION = BUCKET + \"/taxifare/saved_model/\"\n",
    "\n",
    "default_args = {\n",
    "  \"owner\": \"airflow\",\n",
    "  \"depends_on_past\": False,\n",
    "  \"start_date\": airflow.utils.dates.days_ago(2),\n",
    "  \"email\": [\"airflow@example.com\"],\n",
    "  \"email_on_failure\": True,\n",
    "  \"email_on_retry\": False,\n",
    "  \"retries\": 5,\n",
    "  \"retry_delay\": datetime.timedelta(minutes=5)\n",
    "}\n",
    "\n",
    "# Default schedule interval using cronjob syntax - can be customized here\n",
    "# or in the Airflow console.\n",
    "\n",
    "# Specify a schedule interval in CRON syntax to run once a day at 2100 hours (9pm)\n",
    "# Reference: https://airflow.apache.org/scheduler.html\n",
    "schedule_interval = \"00 21 * * *\"\n",
    "\n",
    "# Title your DAG\n",
    "dag = DAG(\n",
    "  \"taxifare_module\", \n",
    "  default_args=default_args,\n",
    "  schedule_interval=None\n",
    ")\n",
    "\n",
    "dag.doc_md = __doc__\n",
    "\n",
    "\n",
    "#\n",
    "#\n",
    "# Task Definition\n",
    "#\n",
    "#\n",
    "\n",
    "for model in SOURCE_DATASET_TABLE_NAMES:\n",
    "  (bq_export_gcs_train_csv_op,\n",
    "   bq_export_gcs_eval_csv_op) = preprocess.preprocess_tasks(\n",
    "       model, dag, PROJECT_ID, BUCKET, DATA_DIR)\n",
    "\n",
    "\n",
    "  (ml_engine_training_op,\n",
    "   bash_copy_new_saved_model_op) = training.training_tasks(\n",
    "       model, dag, PROJECT_ID, BUCKET, DATA_DIR, MODEL_NAME, MODEL_VERSION, MODEL_LOCATION)\n",
    "\n",
    "  (bash_ml_engine_models_list_op,\n",
    "   check_if_model_version_already_exists_op,\n",
    "   ml_engine_create_version_op,\n",
    "   ml_engine_create_other_version_op) = deploy.deploy_tasks(\n",
    "       model, dag, PROJECT_ID, MODEL_NAME, MODEL_VERSION, MODEL_LOCATION)\n",
    "\n",
    "  # Build dependency graph, set_upstream dependencies for all tasks\n",
    "  ml_engine_training_op.set_upstream([bq_export_gcs_train_csv_op, bq_export_gcs_eval_csv_op])\n",
    "\n",
    "  bash_ml_engine_models_list_op.set_upstream(ml_engine_training_op)\n",
    "\n",
    "  ml_engine_create_version_op.set_upstream([bash_copy_new_saved_model_op, check_if_model_version_already_exists_op])\n",
    "  ml_engine_create_other_version_op.set_upstream([bash_copy_new_saved_model_op, check_if_model_version_already_exists_op])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile airflow/dags/module/preprocess.py\n",
    "from airflow.contrib.operators.bigquery_operator import BigQueryOperator\n",
    "from airflow.contrib.operators.bigquery_check_operator import BigQueryCheckOperator\n",
    "from airflow.contrib.operators.bigquery_to_gcs import BigQueryToCloudStorageOperator\n",
    "from airflow.operators.bash_operator import BashOperator\n",
    "\n",
    "\n",
    "def preprocess_tasks(model, dag, PROJECT_ID, BUCKET, DATA_DIR):\n",
    "  # Constants\n",
    "  # Specify your source BigQuery project, dataset, and table names\n",
    "  SOURCE_BQ_PROJECT = \"nyc-tlc\"\n",
    "  SOURCE_DATASET_TABLE_NAMES = \"yellow.trips,green.trips_2014,green.trips_2015\".split(\",\")\n",
    "\n",
    "  # Specify your destination BigQuery dataset\n",
    "  DESTINATION_DATASET = \"taxifare\"\n",
    "  \n",
    "  # BigQuery data query\n",
    "  bql=\"\"\"\n",
    "  SELECT\n",
    "    (tolls_amount + fare_amount) AS fare_amount,\n",
    "    EXTRACT(DAYOFWEEK FROM pickup_datetime) * 1.0 AS dayofweek,\n",
    "    EXTRACT(HOUR FROM pickup_datetime) * 1.0 AS hourofday,\n",
    "    pickup_longitude AS pickuplon,\n",
    "    pickup_latitude AS pickuplat,\n",
    "    dropoff_longitude AS dropofflon,\n",
    "    dropoff_latitude AS dropofflat,\n",
    "    passenger_count*1.0 AS passengers,\n",
    "    CONCAT(CAST(pickup_datetime AS STRING), CAST(pickup_longitude AS STRING), CAST(pickup_latitude AS STRING), CAST(dropoff_latitude AS STRING), CAST(dropoff_longitude AS STRING)) AS key\n",
    "  FROM\n",
    "    `{0}.{1}`\n",
    "  WHERE\n",
    "    trip_distance > 0\n",
    "    AND fare_amount >= 2.5\n",
    "    AND pickup_longitude > -78\n",
    "    AND pickup_longitude < -70\n",
    "    AND dropoff_longitude > -78\n",
    "    AND dropoff_longitude < -70\n",
    "    AND pickup_latitude > 37\n",
    "    AND pickup_latitude < 45\n",
    "    AND dropoff_latitude > 37\n",
    "    AND dropoff_latitude < 45\n",
    "    AND passenger_count > 0\n",
    "    AND rand() < 0.00001\n",
    "  \"\"\"\n",
    "\n",
    "  bql = bql.format(SOURCE_BQ_PROJECT, model)\n",
    "\n",
    "  bql_train = \"SELECT * EXCEPT (key) FROM({0}) WHERE ABS(MOD(FARM_FINGERPRINT(key), 5)) < 4\".format(bql)\n",
    "  bql_eval = \"SELECT * EXCEPT (key) FROM({0}) WHERE ABS(MOD(FARM_FINGERPRINT(key), 5)) = 4\".format(bql)\n",
    "\n",
    "  # Complete the BigQueryOperator task to truncate the table if it already exists before writing\n",
    "  # Reference: https://airflow.apache.org/integration.html#bigqueryoperator\n",
    "  bq_train_data_op = BigQueryOperator(\n",
    "    task_id=\"bq_train_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bql=bql_train,\n",
    "    destination_dataset_table=\"{}.{}_train_data\".format(DESTINATION_DATASET, model.replace(\".\",\"_\")),\n",
    "    write_disposition=\"WRITE_TRUNCATE\", # specify to truncate on writes\n",
    "    use_legacy_sql=False,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bq_eval_data_op = BigQueryOperator(\n",
    "    task_id=\"bq_eval_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bql=bql_eval,\n",
    "    destination_dataset_table=\"{}.{}_eval_data\".format(DESTINATION_DATASET, model.replace(\".\",\"_\")),\n",
    "    write_disposition=\"WRITE_TRUNCATE\", # specify to truncate on writes\n",
    "    use_legacy_sql=False,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  sql = \"\"\"\n",
    "  SELECT\n",
    "    COUNT(*)\n",
    "  FROM\n",
    "    [{0}:{1}.{2}]\n",
    "  \"\"\"\n",
    "\n",
    "  # Check to make sure that the data tables won\"t be empty\n",
    "  bq_check_train_data_op = BigQueryCheckOperator(\n",
    "    task_id=\"bq_check_train_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    sql=sql.format(PROJECT_ID, DESTINATION_DATASET, model.replace(\".\",\"_\") + \"_train_data\"),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bq_check_eval_data_op = BigQueryCheckOperator(\n",
    "    task_id=\"bq_check_eval_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    sql=sql.format(PROJECT_ID, DESTINATION_DATASET, model.replace(\".\",\"_\") + \"_eval_data\"),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # BigQuery training data export to GCS\n",
    "  bash_remove_old_data_op = BashOperator(\n",
    "    task_id=\"bash_remove_old_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bash_command=\"if gsutil ls {0}/taxifare/data/{1} 2> /dev/null; then gsutil -m rm -rf {0}/taxifare/data/{1}/*; else true; fi\".format(BUCKET, model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # Takes a BigQuery dataset and table as input and exports it to GCS as a CSV\n",
    "  bq_export_gcs_train_csv_op = BigQueryToCloudStorageOperator(\n",
    "    task_id=\"bq_export_gcs_train_csv_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    source_project_dataset_table=\"{}.{}_train_data\".format(DESTINATION_DATASET, model.replace(\".\",\"_\")),\n",
    "    destination_cloud_storage_uris=[DATA_DIR + \"{}/train-*.csv\".format(model.replace(\".\",\"_\"))],\n",
    "    export_format=\"CSV\",\n",
    "    print_header=False,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bq_export_gcs_eval_csv_op = BigQueryToCloudStorageOperator(\n",
    "    task_id=\"bq_export_gcs_eval_csv_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    source_project_dataset_table=\"{}.{}_eval_data\".format(DESTINATION_DATASET, model.replace(\".\",\"_\")),\n",
    "    destination_cloud_storage_uris=[DATA_DIR + \"{}/eval-*.csv\".format(model.replace(\".\",\"_\"))],\n",
    "    export_format=\"CSV\",\n",
    "    print_header=False,\n",
    "    dag=dag\n",
    "  )\n",
    "  \n",
    "  # Build dependency graph, set_upstream dependencies for all tasks\n",
    "  bq_check_train_data_op.set_upstream(bq_train_data_op)\n",
    "  bq_check_eval_data_op.set_upstream(bq_eval_data_op)\n",
    "\n",
    "  bash_remove_old_data_op.set_upstream([bq_check_train_data_op, bq_check_eval_data_op])\n",
    "\n",
    "  bq_export_gcs_train_csv_op.set_upstream(bash_remove_old_data_op)\n",
    "  bq_export_gcs_eval_csv_op.set_upstream(bash_remove_old_data_op)\n",
    "  \n",
    "  return (bq_export_gcs_train_csv_op,\n",
    "          bq_export_gcs_eval_csv_op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile airflow/dags/module/training.py\n",
    "import datetime\n",
    "\n",
    "from airflow.contrib.operators.mlengine_operator import MLEngineTrainingOperator\n",
    "from airflow.operators.bash_operator import BashOperator\n",
    "\n",
    "\n",
    "def training_tasks(model, dag, PROJECT_ID, BUCKET, DATA_DIR, MODEL_NAME, MODEL_VERSION, MODEL_LOCATION):\n",
    "  # Constants\n",
    "  # The code package name comes from the model code in the module directory\n",
    "  REGION = \"us-east1\"\n",
    "  PACKAGE_URI = BUCKET + \"/taxifare/code/taxifare-0.1.tar.gz\"\n",
    "  JOB_DIR = BUCKET + \"/jobs\"\n",
    "\n",
    "  # ML Engine training job\n",
    "  job_id = \"taxifare_{}_{}\".format(model.replace(\".\",\"_\"), datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\"))\n",
    "  train_files = DATA_DIR + \"{}/train-*.csv\".format(model.replace(\".\",\"_\"))\n",
    "  eval_files = DATA_DIR + \"{}/eval-*.csv\".format(model.replace(\".\",\"_\"))\n",
    "  output_dir = BUCKET + \"/taxifare/trained_model/{}\".format(model.replace(\".\",\"_\"))\n",
    "  job_dir = JOB_DIR + \"/\" + job_id\n",
    "  training_args = [\n",
    "    \"--job-dir\", job_dir,\n",
    "    \"--train_data_paths\", train_files,\n",
    "    \"--eval_data_paths\", eval_files,\n",
    "    \"--output_dir\", output_dir,\n",
    "    \"--train_steps\", str(500),\n",
    "    \"--train_batch_size\", str(32),\n",
    "    \"--eval_steps\", str(500),\n",
    "    \"--eval_batch_size\", str(32),\n",
    "    \"--nbuckets\", str(8),\n",
    "    \"--hidden_units\", \"128,32,4\"\n",
    "  ]\n",
    "\n",
    "  # Reference: https://airflow.apache.org/integration.html#cloud-ml-engine\n",
    "  ml_engine_training_op = MLEngineTrainingOperator(\n",
    "    task_id=\"ml_engine_training_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID,\n",
    "    job_id=job_id,\n",
    "    package_uris=[PACKAGE_URI],\n",
    "    training_python_module=\"trainer.task\",\n",
    "    training_args=training_args,\n",
    "    region=REGION,\n",
    "    scale_tier=\"BASIC\",\n",
    "    runtime_version=\"1.13\", \n",
    "    python_version=\"3.5\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bash_remove_old_saved_model_op = BashOperator(\n",
    "    task_id=\"bash_remove_old_saved_model_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bash_command=\"if gsutil ls {0} 2> /dev/null; then gsutil -m rm -rf {0}/*; else true; fi\".format(MODEL_LOCATION + model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bash_copy_new_saved_model_op = BashOperator(\n",
    "    task_id=\"bash_copy_new_saved_model_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bash_command=\"gsutil -m rsync -d -r `gsutil ls {0}/export/exporter/ | tail -1` {1}\".format(output_dir, MODEL_LOCATION + model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "  \n",
    "  # Build dependency graph, set_upstream dependencies for all tasks\n",
    "  bash_remove_old_saved_model_op.set_upstream(ml_engine_training_op)\n",
    "  bash_copy_new_saved_model_op.set_upstream(bash_remove_old_saved_model_op)\n",
    "  \n",
    "  return (ml_engine_training_op,\n",
    "          bash_copy_new_saved_model_op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile airflow/dags/module/deploy.py\n",
    "import datetime\n",
    "import logging\n",
    "\n",
    "from airflow.operators.bash_operator import BashOperator\n",
    "from airflow.operators.python_operator import BranchPythonOperator\n",
    "from airflow.operators.dummy_operator import DummyOperator\n",
    "from airflow.contrib.operators.mlengine_operator import MLEngineModelOperator, MLEngineVersionOperator\n",
    "\n",
    "\n",
    "def deploy_tasks(model, dag, PROJECT_ID, MODEL_NAME, MODEL_VERSION, MODEL_LOCATION):\n",
    "  # Constants\n",
    "  OTHER_VERSION_NAME = \"v_{0}\".format(datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")[0:12])\n",
    "\n",
    "  # Create model on ML-Engine\n",
    "  bash_ml_engine_models_list_op = BashOperator(\n",
    "    task_id=\"bash_ml_engine_models_list_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    xcom_push=True,\n",
    "    bash_command=\"gcloud ml-engine models list --filter='name:{0}'\".format(MODEL_NAME + model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  def check_if_model_already_exists(templates_dict, **kwargs):\n",
    "    cur_model = templates_dict[\"model\"].replace(\".\",\"_\")\n",
    "    ml_engine_models_list = kwargs[\"ti\"].xcom_pull(task_ids=\"bash_ml_engine_models_list_{}_task\".format(cur_model))\n",
    "    logging.info(\"check_if_model_already_exists: {}: ml_engine_models_list = \\n{}\".format(cur_model, ml_engine_models_list))\n",
    "    create_model_task = \"ml_engine_create_model_{}_task\".format(cur_model)\n",
    "    dont_create_model_task = \"dont_create_model_dummy_branch_{}_task\".format(cur_model)\n",
    "    if len(ml_engine_models_list) == 0 or ml_engine_models_list == \"Listed 0 items.\":\n",
    "      return create_model_task\n",
    "    return dont_create_model_task\n",
    "\n",
    "  check_if_model_already_exists_op = BranchPythonOperator(\n",
    "    task_id=\"check_if_model_already_exists_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    templates_dict={\"model\": model.replace(\".\",\"_\")},\n",
    "    python_callable=check_if_model_already_exists,\n",
    "    provide_context=True,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_create_model_op = MLEngineModelOperator(\n",
    "    task_id=\"ml_engine_create_model_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model={\"name\": MODEL_NAME + model.replace(\".\",\"_\")}, \n",
    "    operation=\"create\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  create_model_dummy_op = DummyOperator(\n",
    "    task_id=\"create_model_dummy_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    trigger_rule=\"all_done\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  dont_create_model_dummy_branch_op = DummyOperator(\n",
    "    task_id=\"dont_create_model_dummy_branch_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  dont_create_model_dummy_op = DummyOperator(\n",
    "    task_id=\"dont_create_model_dummy_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    trigger_rule=\"all_done\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # Create version of model on ML-Engine\n",
    "  bash_ml_engine_versions_list_op = BashOperator(\n",
    "    task_id=\"bash_ml_engine_versions_list_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    xcom_push=True,\n",
    "    bash_command=\"gcloud ml-engine versions list --model {0} --filter='name:{1}'\".format(MODEL_NAME + model.replace(\".\",\"_\"), MODEL_VERSION),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  def check_if_model_version_already_exists(templates_dict, **kwargs):\n",
    "    cur_model = templates_dict[\"model\"].replace(\".\",\"_\")\n",
    "    ml_engine_versions_list = kwargs[\"ti\"].xcom_pull(task_ids=\"bash_ml_engine_versions_list_{}_task\".format(cur_model))\n",
    "    logging.info(\"check_if_model_version_already_exists: {}: ml_engine_versions_list = \\n{}\".format(cur_model, ml_engine_versions_list))\n",
    "    create_version_task = \"ml_engine_create_version_{}_task\".format(cur_model)\n",
    "    create_other_version_task = \"ml_engine_create_other_version_{}_task\".format(cur_model)\n",
    "    if len(ml_engine_versions_list) == 0 or ml_engine_versions_list == \"Listed 0 items.\":\n",
    "      return create_version_task\n",
    "    return create_other_version_task\n",
    "\n",
    "  check_if_model_version_already_exists_op = BranchPythonOperator(\n",
    "    task_id=\"check_if_model_version_already_exists_{}_task\".format(model.replace(\".\",\"_\")), \n",
    "    templates_dict={\"model\": model.replace(\".\",\"_\")},\n",
    "    python_callable=check_if_model_version_already_exists,\n",
    "    provide_context=True,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_create_version_op = MLEngineVersionOperator(\n",
    "    task_id=\"ml_engine_create_version_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model_name=MODEL_NAME + model.replace(\".\",\"_\"), \n",
    "    version_name=MODEL_VERSION, \n",
    "    version={\n",
    "      \"name\": MODEL_VERSION,\n",
    "      \"deploymentUri\": MODEL_LOCATION + model.replace(\".\",\"_\"),\n",
    "      \"runtimeVersion\": \"1.13\",\n",
    "      \"framework\": \"TENSORFLOW\",\n",
    "      \"pythonVersion\": \"3.5\",\n",
    "    },\n",
    "    operation=\"create\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_create_other_version_op = MLEngineVersionOperator(\n",
    "    task_id=\"ml_engine_create_other_version_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model_name=MODEL_NAME + model.replace(\".\",\"_\"), \n",
    "    version_name=OTHER_VERSION_NAME, \n",
    "    version={\n",
    "      \"name\": OTHER_VERSION_NAME,\n",
    "      \"deploymentUri\": MODEL_LOCATION + model.replace(\".\",\"_\"),\n",
    "      \"runtimeVersion\": \"1.13\",\n",
    "      \"framework\": \"TENSORFLOW\",\n",
    "      \"pythonVersion\": \"3.5\",\n",
    "    },\n",
    "    operation=\"create\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_set_default_version_op = MLEngineVersionOperator(\n",
    "    task_id=\"ml_engine_set_default_version_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model_name=MODEL_NAME + model.replace(\".\",\"_\"), \n",
    "    version_name=MODEL_VERSION, \n",
    "    version={\"name\": MODEL_VERSION}, \n",
    "    operation=\"set_default\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_set_default_other_version_op = MLEngineVersionOperator(\n",
    "    task_id=\"ml_engine_set_default_other_version_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model_name=MODEL_NAME + model.replace(\".\",\"_\"), \n",
    "    version_name=OTHER_VERSION_NAME, \n",
    "    version={\"name\": OTHER_VERSION_NAME}, \n",
    "    operation=\"set_default\",\n",
    "    dag=dag\n",
    "  )\n",
    "  \n",
    "  # Build dependency graph, set_upstream dependencies for all tasks\n",
    "  check_if_model_already_exists_op.set_upstream(bash_ml_engine_models_list_op)\n",
    "\n",
    "  ml_engine_create_model_op.set_upstream(check_if_model_already_exists_op)\n",
    "  create_model_dummy_op.set_upstream(ml_engine_create_model_op)\n",
    "  dont_create_model_dummy_branch_op.set_upstream(check_if_model_already_exists_op)\n",
    "  dont_create_model_dummy_op.set_upstream(dont_create_model_dummy_branch_op)\n",
    "\n",
    "  bash_ml_engine_versions_list_op.set_upstream([dont_create_model_dummy_op, create_model_dummy_op])\n",
    "  check_if_model_version_already_exists_op.set_upstream(bash_ml_engine_versions_list_op)\n",
    "\n",
    "  ml_engine_set_default_version_op.set_upstream(ml_engine_create_version_op)\n",
    "  ml_engine_set_default_other_version_op.set_upstream(ml_engine_create_other_version_op)\n",
    "  \n",
    "  return (bash_ml_engine_models_list_op,\n",
    "          check_if_model_version_already_exists_op,\n",
    "          ml_engine_create_version_op,\n",
    "          ml_engine_create_other_version_op)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Subdag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile airflow/dags/subdag/taxifare_subdag.py\n",
    "# Copyright 2018 Google Inc. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "\n",
    "\"\"\"DAG definition for taxifare automated pipeline.\"\"\"\n",
    "\n",
    "import airflow\n",
    "from airflow import DAG\n",
    "\n",
    "# Reference for all available airflow operators: \n",
    "# https://github.com/apache/incubator-airflow/tree/master/airflow/contrib/operators\n",
    "from airflow.operators.subdag_operator import SubDagOperator\n",
    "from airflow.hooks.base_hook import BaseHook\n",
    "\n",
    "from airflow.models import TaskInstance\n",
    "\n",
    "import datetime\n",
    "\n",
    "from subdag import preprocess\n",
    "from subdag import training\n",
    "from subdag import deploy\n",
    "\n",
    "\n",
    "def _get_project_id():\n",
    "  \"\"\"Get project ID from default GCP connection.\"\"\"\n",
    "\n",
    "  extras = BaseHook.get_connection(\"google_cloud_default\").extra_dejson\n",
    "  key = \"extra__google_cloud_platform__project\"\n",
    "  if key in extras:\n",
    "    project_id = extras[key]\n",
    "  else:\n",
    "    raise (\"Must configure project_id in google_cloud_default \"\n",
    "           \"connection from Airflow Console\")\n",
    "  return project_id\n",
    "\n",
    "# Constants\n",
    "# Get project ID and GCS bucket\n",
    "PROJECT_ID = _get_project_id()\n",
    "BUCKET = \"gs://\" + PROJECT_ID + \"-bucket\"\n",
    "\n",
    "# Specify your source BigQuery dataset and table names\n",
    "SOURCE_DATASET_TABLE_NAMES = \"yellow.trips,green.trips_2014,green.trips_2015\".split(\",\")\n",
    "\n",
    "# Where to write out data in GCS\n",
    "DATA_DIR = BUCKET + \"/taxifare/data/\"\n",
    "\n",
    "# Base model parameters\n",
    "MODEL_NAME = \"taxifare_\"\n",
    "MODEL_VERSION = \"v1\"\n",
    "MODEL_LOCATION = BUCKET + \"/taxifare/saved_model/\"\n",
    "\n",
    "default_args = {\n",
    "  \"owner\": \"airflow\",\n",
    "  \"depends_on_past\": False,\n",
    "  \"start_date\": airflow.utils.dates.days_ago(2),\n",
    "  \"email\": [\"airflow@example.com\"],\n",
    "  \"email_on_failure\": True,\n",
    "  \"email_on_retry\": False,\n",
    "  \"retries\": 5,\n",
    "  \"retry_delay\": datetime.timedelta(minutes=5)\n",
    "}\n",
    "\n",
    "# Default schedule interval using cronjob syntax - can be customized here\n",
    "# or in the Airflow console.\n",
    "\n",
    "# Specify a schedule interval in CRON syntax to run once a day at 2100 hours (9pm)\n",
    "# Reference: https://airflow.apache.org/scheduler.html\n",
    "schedule_interval = \"00 21 * * *\"\n",
    "\n",
    "# Title your DAG\n",
    "DAG_NAME = \"taxifare_subdag\"\n",
    "\n",
    "dag = DAG(\n",
    "  DAG_NAME, \n",
    "  default_args=default_args,\n",
    "  schedule_interval=None\n",
    ")\n",
    "\n",
    "dag.doc_md = __doc__\n",
    "\n",
    "\n",
    "#\n",
    "#\n",
    "# Task Definition\n",
    "#\n",
    "#\n",
    "\n",
    "for model in SOURCE_DATASET_TABLE_NAMES:\n",
    "  subdag_preprocess_op = SubDagOperator(\n",
    "    task_id=\"subdag_preprocess_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    subdag=preprocess.preprocess_tasks(\n",
    "        model,\n",
    "        DAG_NAME,\n",
    "        \"subdag_preprocess_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "        default_args,\n",
    "        PROJECT_ID,\n",
    "        BUCKET,\n",
    "        DATA_DIR),\n",
    "    dag=dag\n",
    "  )\n",
    "  \n",
    "  subdag_training_op = SubDagOperator(\n",
    "    task_id=\"subdag_training_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    subdag=training.training_tasks(\n",
    "        model,\n",
    "        DAG_NAME,\n",
    "        \"subdag_training_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "        default_args,\n",
    "        PROJECT_ID,\n",
    "        BUCKET,\n",
    "        DATA_DIR,\n",
    "        MODEL_NAME,\n",
    "        MODEL_VERSION,\n",
    "        MODEL_LOCATION),\n",
    "    dag=dag\n",
    "  )\n",
    "  \n",
    "  subdag_deploy_op = SubDagOperator(\n",
    "    task_id=\"subdag_deploy_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    subdag=deploy.deploy_tasks(\n",
    "        model,\n",
    "        DAG_NAME,\n",
    "        \"subdag_deploy_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "        default_args,\n",
    "        PROJECT_ID,\n",
    "        MODEL_NAME,\n",
    "        MODEL_VERSION,\n",
    "        MODEL_LOCATION),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # Build dependency graph, set_upstream dependencies for all tasks\n",
    "  subdag_training_op.set_upstream(subdag_preprocess_op)\n",
    "  subdag_deploy_op.set_upstream(subdag_training_op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile airflow/dags/subdag/preprocess.py\n",
    "from airflow import DAG\n",
    "from airflow.contrib.operators.bigquery_operator import BigQueryOperator\n",
    "from airflow.contrib.operators.bigquery_check_operator import BigQueryCheckOperator\n",
    "from airflow.contrib.operators.bigquery_to_gcs import BigQueryToCloudStorageOperator\n",
    "from airflow.operators.bash_operator import BashOperator\n",
    "\n",
    "\n",
    "def preprocess_tasks(model, parent_dag_name, child_dag_name, default_args, PROJECT_ID, BUCKET, DATA_DIR):\n",
    "  # Create inner dag\n",
    "  dag = DAG(\n",
    "    \"{0}.{1}\".format(parent_dag_name, child_dag_name),\n",
    "    default_args=default_args,\n",
    "    schedule_interval=None\n",
    "  )\n",
    "\n",
    "  # Constants\n",
    "  # Specify your source BigQuery project, dataset, and table names\n",
    "  SOURCE_BQ_PROJECT = \"nyc-tlc\"\n",
    "  SOURCE_DATASET_TABLE_NAMES = \"yellow.trips,green.trips_2014,green.trips_2015\".split(\",\")\n",
    "\n",
    "  # Specify your destination BigQuery dataset\n",
    "  DESTINATION_DATASET = \"taxifare\"\n",
    "  \n",
    "  # BigQuery data query\n",
    "  bql=\"\"\"\n",
    "  SELECT\n",
    "    (tolls_amount + fare_amount) AS fare_amount,\n",
    "    EXTRACT(DAYOFWEEK FROM pickup_datetime) * 1.0 AS dayofweek,\n",
    "    EXTRACT(HOUR FROM pickup_datetime) * 1.0 AS hourofday,\n",
    "    pickup_longitude AS pickuplon,\n",
    "    pickup_latitude AS pickuplat,\n",
    "    dropoff_longitude AS dropofflon,\n",
    "    dropoff_latitude AS dropofflat,\n",
    "    passenger_count*1.0 AS passengers,\n",
    "    CONCAT(CAST(pickup_datetime AS STRING), CAST(pickup_longitude AS STRING), CAST(pickup_latitude AS STRING), CAST(dropoff_latitude AS STRING), CAST(dropoff_longitude AS STRING)) AS key\n",
    "  FROM\n",
    "    `{0}.{1}`\n",
    "  WHERE\n",
    "    trip_distance > 0\n",
    "    AND fare_amount >= 2.5\n",
    "    AND pickup_longitude > -78\n",
    "    AND pickup_longitude < -70\n",
    "    AND dropoff_longitude > -78\n",
    "    AND dropoff_longitude < -70\n",
    "    AND pickup_latitude > 37\n",
    "    AND pickup_latitude < 45\n",
    "    AND dropoff_latitude > 37\n",
    "    AND dropoff_latitude < 45\n",
    "    AND passenger_count > 0\n",
    "    AND rand() < 0.00001\n",
    "  \"\"\"\n",
    "\n",
    "  bql = bql.format(SOURCE_BQ_PROJECT, model)\n",
    "\n",
    "  bql_train = \"SELECT * EXCEPT (key) FROM({0}) WHERE ABS(MOD(FARM_FINGERPRINT(key), 5)) < 4\".format(bql)\n",
    "  bql_eval = \"SELECT * EXCEPT (key) FROM({0}) WHERE ABS(MOD(FARM_FINGERPRINT(key), 5)) = 4\".format(bql)\n",
    "\n",
    "  # Complete the BigQueryOperator task to truncate the table if it already exists before writing\n",
    "  # Reference: https://airflow.apache.org/integration.html#bigqueryoperator\n",
    "  bq_train_data_op = BigQueryOperator(\n",
    "    task_id=\"bq_train_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bql=bql_train,\n",
    "    destination_dataset_table=\"{}.{}_train_data\".format(DESTINATION_DATASET, model.replace(\".\",\"_\")),\n",
    "    write_disposition=\"WRITE_TRUNCATE\", # specify to truncate on writes\n",
    "    use_legacy_sql=False,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bq_eval_data_op = BigQueryOperator(\n",
    "    task_id=\"bq_eval_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bql=bql_eval,\n",
    "    destination_dataset_table=\"{}.{}_eval_data\".format(DESTINATION_DATASET, model.replace(\".\",\"_\")),\n",
    "    write_disposition=\"WRITE_TRUNCATE\", # specify to truncate on writes\n",
    "    use_legacy_sql=False,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  sql = \"\"\"\n",
    "  SELECT\n",
    "    COUNT(*)\n",
    "  FROM\n",
    "    [{0}:{1}.{2}]\n",
    "  \"\"\"\n",
    "\n",
    "  # Check to make sure that the data tables won\"t be empty\n",
    "  bq_check_train_data_op = BigQueryCheckOperator(\n",
    "    task_id=\"bq_check_train_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    sql=sql.format(PROJECT_ID, DESTINATION_DATASET, model.replace(\".\",\"_\") + \"_train_data\"),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bq_check_eval_data_op = BigQueryCheckOperator(\n",
    "    task_id=\"bq_check_eval_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    sql=sql.format(PROJECT_ID, DESTINATION_DATASET, model.replace(\".\",\"_\") + \"_eval_data\"),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # BigQuery training data export to GCS\n",
    "  bash_remove_old_data_op = BashOperator(\n",
    "    task_id=\"bash_remove_old_data_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bash_command=\"if gsutil ls {0}/taxifare/data/{1} 2> /dev/null; then gsutil -m rm -rf {0}/taxifare/data/{1}/*; else true; fi\".format(BUCKET, model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # Takes a BigQuery dataset and table as input and exports it to GCS as a CSV\n",
    "  bq_export_gcs_train_csv_op = BigQueryToCloudStorageOperator(\n",
    "    task_id=\"bq_export_gcs_train_csv_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    source_project_dataset_table=\"{}.{}_train_data\".format(DESTINATION_DATASET, model.replace(\".\",\"_\")),\n",
    "    destination_cloud_storage_uris=[DATA_DIR + \"{}/train-*.csv\".format(model.replace(\".\",\"_\"))],\n",
    "    export_format=\"CSV\",\n",
    "    print_header=False,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bq_export_gcs_eval_csv_op = BigQueryToCloudStorageOperator(\n",
    "    task_id=\"bq_export_gcs_eval_csv_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    source_project_dataset_table=\"{}.{}_eval_data\".format(DESTINATION_DATASET, model.replace(\".\",\"_\")),\n",
    "    destination_cloud_storage_uris=[DATA_DIR + \"{}/eval-*.csv\".format(model.replace(\".\",\"_\"))],\n",
    "    export_format=\"CSV\",\n",
    "    print_header=False,\n",
    "    dag=dag\n",
    "  )\n",
    "  \n",
    "  # Build dependency graph, set_upstream dependencies for all tasks\n",
    "  bq_check_train_data_op.set_upstream(bq_train_data_op)\n",
    "  bq_check_eval_data_op.set_upstream(bq_eval_data_op)\n",
    "\n",
    "  bash_remove_old_data_op.set_upstream([bq_check_train_data_op, bq_check_eval_data_op])\n",
    "\n",
    "  bq_export_gcs_train_csv_op.set_upstream(bash_remove_old_data_op)\n",
    "  bq_export_gcs_eval_csv_op.set_upstream(bash_remove_old_data_op)\n",
    "  \n",
    "  return dag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile airflow/dags/subdag/training.py\n",
    "import datetime\n",
    "\n",
    "from airflow import DAG\n",
    "from airflow.contrib.operators.mlengine_operator import MLEngineTrainingOperator\n",
    "from airflow.operators.bash_operator import BashOperator\n",
    "\n",
    "\n",
    "def training_tasks(model, parent_dag_name, child_dag_name, default_args, PROJECT_ID, BUCKET, DATA_DIR, MODEL_NAME, MODEL_VERSION, MODEL_LOCATION):\n",
    "  # Create inner dag\n",
    "  dag = DAG(\n",
    "    \"{0}.{1}\".format(parent_dag_name, child_dag_name),\n",
    "    default_args=default_args,\n",
    "    schedule_interval=None\n",
    "  )\n",
    "\n",
    "  # Constants\n",
    "  # The code package name comes from the model code in the module directory\n",
    "  REGION = \"us-east1\"\n",
    "  PACKAGE_URI = BUCKET + \"/taxifare/code/taxifare-0.1.tar.gz\"\n",
    "  JOB_DIR = BUCKET + \"/jobs\"\n",
    "\n",
    "  # ML Engine training job\n",
    "  job_id = \"taxifare_{}_{}\".format(model.replace(\".\",\"_\"), datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\"))\n",
    "  train_files = DATA_DIR + \"{}/train-*.csv\".format(model.replace(\".\",\"_\"))\n",
    "  eval_files = DATA_DIR + \"{}/eval-*.csv\".format(model.replace(\".\",\"_\"))\n",
    "  output_dir = BUCKET + \"/taxifare/trained_model/{}\".format(model.replace(\".\",\"_\"))\n",
    "  job_dir = JOB_DIR + \"/\" + job_id\n",
    "  training_args = [\n",
    "    \"--job-dir\", job_dir,\n",
    "    \"--train_data_paths\", train_files,\n",
    "    \"--eval_data_paths\", eval_files,\n",
    "    \"--output_dir\", output_dir,\n",
    "    \"--train_steps\", str(500),\n",
    "    \"--train_batch_size\", str(32),\n",
    "    \"--eval_steps\", str(500),\n",
    "    \"--eval_batch_size\", str(32),\n",
    "    \"--nbuckets\", str(8),\n",
    "    \"--hidden_units\", \"128,32,4\"\n",
    "  ]\n",
    "\n",
    "  # Reference: https://airflow.apache.org/integration.html#cloud-ml-engine\n",
    "  ml_engine_training_op = MLEngineTrainingOperator(\n",
    "    task_id=\"ml_engine_training_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID,\n",
    "    job_id=job_id,\n",
    "    package_uris=[PACKAGE_URI],\n",
    "    training_python_module=\"trainer.task\",\n",
    "    training_args=training_args,\n",
    "    region=REGION,\n",
    "    scale_tier=\"BASIC\",\n",
    "    runtime_version=\"1.13\", \n",
    "    python_version=\"3.5\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bash_remove_old_saved_model_op = BashOperator(\n",
    "    task_id=\"bash_remove_old_saved_model_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bash_command=\"if gsutil ls {0} 2> /dev/null; then gsutil -m rm -rf {0}/*; else true; fi\".format(MODEL_LOCATION + model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  bash_copy_new_saved_model_op = BashOperator(\n",
    "    task_id=\"bash_copy_new_saved_model_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    bash_command=\"gsutil -m rsync -d -r `gsutil ls {0}/export/exporter/ | tail -1` {1}\".format(output_dir, MODEL_LOCATION + model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # Build dependency graph, set_upstream dependencies for all tasks\n",
    "  bash_remove_old_saved_model_op.set_upstream(ml_engine_training_op)\n",
    "  bash_copy_new_saved_model_op.set_upstream(bash_remove_old_saved_model_op)\n",
    "  \n",
    "  return dag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile airflow/dags/subdag/deploy.py\n",
    "import datetime\n",
    "import logging\n",
    "\n",
    "from airflow import DAG\n",
    "from airflow.operators.bash_operator import BashOperator\n",
    "from airflow.operators.python_operator import BranchPythonOperator\n",
    "from airflow.operators.dummy_operator import DummyOperator\n",
    "from airflow.contrib.operators.mlengine_operator import MLEngineModelOperator, MLEngineVersionOperator\n",
    "\n",
    "\n",
    "def deploy_tasks(model, parent_dag_name, child_dag_name, default_args, PROJECT_ID, MODEL_NAME, MODEL_VERSION, MODEL_LOCATION):\n",
    "  # Create inner dag\n",
    "  dag = DAG(\n",
    "    \"{0}.{1}\".format(parent_dag_name, child_dag_name),\n",
    "    default_args=default_args,\n",
    "    schedule_interval=None\n",
    "  )\n",
    "\n",
    "  # Constants\n",
    "  OTHER_VERSION_NAME = \"v_{0}\".format(datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")[0:12])\n",
    "\n",
    "  # Create model on ML-Engine\n",
    "  bash_ml_engine_models_list_op = BashOperator(\n",
    "    task_id=\"bash_ml_engine_models_list_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    xcom_push=True,\n",
    "    bash_command=\"gcloud ml-engine models list --filter='name:{0}'\".format(MODEL_NAME + model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  def check_if_model_already_exists(templates_dict, **kwargs):\n",
    "    cur_model = templates_dict[\"model\"].replace(\".\",\"_\")\n",
    "    ml_engine_models_list = kwargs[\"ti\"].xcom_pull(task_ids=\"bash_ml_engine_models_list_{}_task\".format(cur_model))\n",
    "    logging.info(\"check_if_model_already_exists: {}: ml_engine_models_list = \\n{}\".format(cur_model, ml_engine_models_list))\n",
    "    create_model_task = \"ml_engine_create_model_{}_task\".format(cur_model)\n",
    "    dont_create_model_task = \"dont_create_model_dummy_branch_{}_task\".format(cur_model)\n",
    "    if len(ml_engine_models_list) == 0 or ml_engine_models_list == \"Listed 0 items.\":\n",
    "      return create_model_task\n",
    "    return dont_create_model_task\n",
    "\n",
    "  check_if_model_already_exists_op = BranchPythonOperator(\n",
    "      task_id=\"check_if_model_already_exists_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "      templates_dict={\"model\": model.replace(\".\",\"_\")},\n",
    "      python_callable=check_if_model_already_exists,\n",
    "      provide_context=True,\n",
    "      dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_create_model_op = MLEngineModelOperator(\n",
    "    task_id=\"ml_engine_create_model_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model={\"name\": MODEL_NAME + model.replace(\".\",\"_\")}, \n",
    "    operation=\"create\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  create_model_dummy_op = DummyOperator(\n",
    "    task_id=\"create_model_dummy_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    trigger_rule=\"all_done\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  dont_create_model_dummy_branch_op = DummyOperator(\n",
    "    task_id=\"dont_create_model_dummy_branch_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  dont_create_model_dummy_op = DummyOperator(\n",
    "    task_id=\"dont_create_model_dummy_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    trigger_rule=\"all_done\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # Create version of model on ML-Engine\n",
    "  bash_ml_engine_versions_list_op = BashOperator(\n",
    "    task_id=\"bash_ml_engine_versions_list_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    xcom_push=True,\n",
    "    bash_command=\"gcloud ml-engine versions list --model {0} --filter='name:{1}'\".format(MODEL_NAME + model.replace(\".\",\"_\"), MODEL_VERSION),\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  def check_if_model_version_already_exists(templates_dict, **kwargs):\n",
    "    cur_model = templates_dict[\"model\"].replace(\".\",\"_\")\n",
    "    ml_engine_versions_list = kwargs[\"ti\"].xcom_pull(task_ids=\"bash_ml_engine_versions_list_{}_task\".format(cur_model))\n",
    "    logging.info(\"check_if_model_version_already_exists: {}: ml_engine_versions_list = \\n{}\".format(cur_model, ml_engine_versions_list))\n",
    "    create_version_task = \"ml_engine_create_version_{}_task\".format(cur_model)\n",
    "    create_other_version_task = \"ml_engine_create_other_version_{}_task\".format(cur_model)\n",
    "    if len(ml_engine_versions_list) == 0 or ml_engine_versions_list == \"Listed 0 items.\":\n",
    "      return create_version_task\n",
    "    return create_other_version_task\n",
    "\n",
    "  check_if_model_version_already_exists_op = BranchPythonOperator(\n",
    "    task_id=\"check_if_model_version_already_exists_{}_task\".format(model.replace(\".\",\"_\")), \n",
    "    templates_dict={\"model\": model.replace(\".\",\"_\")},\n",
    "    python_callable=check_if_model_version_already_exists,\n",
    "    provide_context=True,\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_create_version_op = MLEngineVersionOperator(\n",
    "    task_id=\"ml_engine_create_version_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model_name=MODEL_NAME + model.replace(\".\",\"_\"), \n",
    "    version_name=MODEL_VERSION, \n",
    "    version={\n",
    "        \"name\": MODEL_VERSION,\n",
    "        \"deploymentUri\": MODEL_LOCATION + model.replace(\".\",\"_\"),\n",
    "        \"runtimeVersion\": \"1.13\",\n",
    "        \"framework\": \"TENSORFLOW\",\n",
    "        \"pythonVersion\": \"3.5\",\n",
    "    },\n",
    "    operation=\"create\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_create_other_version_op = MLEngineVersionOperator(\n",
    "    task_id=\"ml_engine_create_other_version_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model_name=MODEL_NAME + model.replace(\".\",\"_\"), \n",
    "    version_name=OTHER_VERSION_NAME, \n",
    "    version={\n",
    "        \"name\": OTHER_VERSION_NAME,\n",
    "        \"deploymentUri\": MODEL_LOCATION + model.replace(\".\",\"_\"),\n",
    "        \"runtimeVersion\": \"1.13\",\n",
    "        \"framework\": \"TENSORFLOW\",\n",
    "        \"pythonVersion\": \"3.5\",\n",
    "    },\n",
    "    operation=\"create\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_set_default_version_op = MLEngineVersionOperator(\n",
    "    task_id=\"ml_engine_set_default_version_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model_name=MODEL_NAME + model.replace(\".\",\"_\"), \n",
    "    version_name=MODEL_VERSION, \n",
    "    version={\"name\": MODEL_VERSION}, \n",
    "    operation=\"set_default\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  ml_engine_set_default_other_version_op = MLEngineVersionOperator(\n",
    "    task_id=\"ml_engine_set_default_other_version_{}_task\".format(model.replace(\".\",\"_\")),\n",
    "    project_id=PROJECT_ID, \n",
    "    model_name=MODEL_NAME + model.replace(\".\",\"_\"), \n",
    "    version_name=OTHER_VERSION_NAME, \n",
    "    version={\"name\": OTHER_VERSION_NAME}, \n",
    "    operation=\"set_default\",\n",
    "    dag=dag\n",
    "  )\n",
    "\n",
    "  # Build dependency graph, set_upstream dependencies for all tasks\n",
    "  check_if_model_already_exists_op.set_upstream(bash_ml_engine_models_list_op)\n",
    "\n",
    "  ml_engine_create_model_op.set_upstream(check_if_model_already_exists_op)\n",
    "  create_model_dummy_op.set_upstream(ml_engine_create_model_op)\n",
    "  dont_create_model_dummy_branch_op.set_upstream(check_if_model_already_exists_op)\n",
    "  dont_create_model_dummy_op.set_upstream(dont_create_model_dummy_branch_op)\n",
    "\n",
    "  bash_ml_engine_versions_list_op.set_upstream([dont_create_model_dummy_op, create_model_dummy_op])\n",
    "  check_if_model_version_already_exists_op.set_upstream(bash_ml_engine_versions_list_op)\n",
    "\n",
    "  ml_engine_create_version_op.set_upstream(check_if_model_version_already_exists_op)\n",
    "  ml_engine_create_other_version_op.set_upstream(check_if_model_version_already_exists_op)\n",
    "\n",
    "  ml_engine_set_default_version_op.set_upstream(ml_engine_create_version_op)\n",
    "  ml_engine_set_default_other_version_op.set_upstream(ml_engine_create_other_version_op)\n",
    "  \n",
    "  return dag"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2-YF4aOUm1Pl"
   },
   "source": [
    "### Copy local Airflow DAG file and plugins into the DAGs folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KAd9I3icm1Pm"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "gsutil -m cp -r airflow/dags/* gs://${AIRFLOW_BUCKET}/dags # overwrite if it exists"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0ixqSbrmm1Po"
   },
   "source": [
    "1. Navigate to your Cloud Composer [instance](https://console.cloud.google.com/composer/environments?project=)<br/><br/>\n",
    "\n",
    "2. Trigger a __manual run__ of your DAG for testing<br/><br/>\n",
    "\n",
    "3. Ensure your DAG runs successfully (all nodes outlined in dark green and 'success' tag shows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "cloud_composer_automated_ml_pipeline_taxifare.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
